package com.cfp4cloud.cfp.knowledge.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.SecureUtil;
import cn.hutool.extra.template.Template;
import cn.hutool.extra.template.TemplateConfig;
import cn.hutool.extra.template.TemplateEngine;
import cn.hutool.extra.template.TemplateUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.cfp4cloud.cfp.knowledge.dto.AiDataTableDTO;
import com.cfp4cloud.cfp.knowledge.entity.AiDataEntity;
import com.cfp4cloud.cfp.knowledge.entity.AiDataTableEntity;
import com.cfp4cloud.cfp.knowledge.mapper.AiDataMapper;
import com.cfp4cloud.cfp.knowledge.mapper.AiDataTableMapper;
import com.cfp4cloud.cfp.knowledge.service.AiDataFieldService;
import com.cfp4cloud.cfp.knowledge.service.AiDataService;
import com.cfp4cloud.cfp.knowledge.support.constant.EmbedBizTypeEnums;
import com.cfp4cloud.cfp.knowledge.support.provider.MemoryEmbeddingProvider;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.segment.TextSegment;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Lazy;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;

import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static com.cfp4cloud.cfp.knowledge.support.provider.MemoryEmbeddingProvider.TEMP_ID;

/**
 * AI数据集管理服务实现类
 * <p>
 * 管理多表数据集的定义和元数据 提供表结构查询和Schema生成功能 支持虚拟注释增强SQL生成质量
 *
 * @author cfp
 * @date 2025-03-26 21:47:45
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class AiDataServiceImpl extends ServiceImpl<AiDataMapper, AiDataEntity> implements AiDataService {

	/**
	 * 模板引擎 用于生成表结构描述文档
	 */
	public static final TemplateEngine engine = TemplateUtil.createEngine(new TemplateConfig());

	/**
	 * 数据表数据访问层
	 */
	private final AiDataTableMapper dataTableMapper;

	/**
	 * 数据字段服务（延迟加载） 避免循环依赖
	 */
	@Lazy
	private final AiDataFieldService dataFieldService;

	/**
	 * 数据集数据访问层
	 */
	private final AiDataMapper aiDataMapper;

	/**
	 * 表结构模板资源 FreeMarker模板文件
	 */
	@Value("classpath:/prompts/chat2db-table-schema.ftl")
	private Resource tableSchema;

	/**
	 * 查询数据集的Schema定义
	 * <p>
	 * 获取数据集包含的所有表结构信息 使用模板渲染成结构化文档 优先使用虚拟注释提高可读性
	 * @param dataId 数据集ID
	 * @return 格式化的Schema描述文档
	 */
	@SneakyThrows
	@Override
	public String queryDataSchema(Long dataId) {
		AiDataEntity aiDataEntity = aiDataMapper.selectById(dataId);
		String[] tableNames = aiDataEntity.getTableName();
		String dsName = aiDataEntity.getDsName();

		// 查询数据集包含的所有表
		List<AiDataTableEntity> tableEntityList = dataTableMapper.selectList(Wrappers.<AiDataTableEntity>lambdaQuery()
			.eq(AiDataTableEntity::getDsName, dsName)
			.in(AiDataTableEntity::getTableName, CollUtil.toList(tableNames)));

		// 渲染表结构文档
		String renderResult = queryTableSchema(tableEntityList);
		log.info("目标数据集，渲染结果：{}", renderResult);
		return renderResult;
	}

	/**
	 * 查询表结构Schema
	 * <p>
	 * 将表和字段信息转换为结构化文档 优先使用AI优化后的虚拟注释
	 * @param tableEntityList 表实体列表
	 * @return 模板渲染后的Schema文档
	 */
	@SneakyThrows
	public String queryTableSchema(List<AiDataTableEntity> tableEntityList) {
		List<AiDataTableDTO> tableDTOList = new ArrayList<>();
		for (AiDataTableEntity tableEntity : tableEntityList) {
			AiDataTableDTO tableDTO = new AiDataTableDTO();
			tableDTO.setTableName(tableEntity.getTableName());
			tableDTO.setDbType(tableEntity.getDsType());
			// 优先使用虚拟注释，否则使用原始注释
			tableDTO.setTableComment(StrUtil.isNotBlank(tableEntity.getVirtualComment())
					? tableEntity.getVirtualComment() : tableEntity.getTableComment());

			List<AiDataTableDTO.AiTableField> fieldList = dataFieldService
				.getOrSync(tableEntity.getDsName(), tableEntity.getTableName())
				.stream()
				.map(fieldEntity -> {
					AiDataTableDTO.AiTableField field = new AiDataTableDTO.AiTableField();
					field.setFieldName(fieldEntity.getFieldName());
					// 字段也优先使用虚拟注释
					field.setFieldComment(StrUtil.isNotBlank(fieldEntity.getVirtualComment())
							? fieldEntity.getVirtualComment() : fieldEntity.getFieldComment());
					field.setDbType(fieldEntity.getDbType());
					field.setFieldType(fieldEntity.getFieldType());
					return field;
				})
				.toList();
			tableDTO.setFields(fieldList);
			tableDTOList.add(tableDTO);
		}

		Template template = engine.getTemplate(IoUtil.read(tableSchema.getInputStream(), Charset.defaultCharset()));
		return template.render(Map.of("tables", tableDTOList));
	}

	/**
	 * 保存或更新数据集
	 * <p>
	 * 同时将数据集描述信息向量化 存储到内存向量库供快速检索 用于Chat2SQL场景的数据集匹配
	 * @param aiData AI数据集实体
	 * @return 保存成功返回true
	 */
	@Override
	public boolean saveOrUpdateData(AiDataEntity aiData) {
		baseMapper.insertOrUpdate(aiData);
		// 创建文本段落，包含数据源和描述信息
		TextSegment textSegment = TextSegment.textSegment(aiData.getDsName() + aiData.getDescription(),
				Metadata.from(Map.of(EmbedBizTypeEnums.Fields.type, EmbedBizTypeEnums.CHAT2SQL.getType(), TEMP_ID,
						aiData.getDataId())));
		// 添加到内存向量存储
		MemoryEmbeddingProvider.add(SecureUtil.md5(aiData.getDsName()), textSegment);
		return true;
	}

}
